import os
import pandas as pd
from PIL import Image
import cv2
import numpy as np



def rust_estimation(img_files_path, label_files_path, f_dx):

    img_files = os.listdir(img_files_path)
    result = []
    for img_name in img_files:
        img_path = os.path.join(img_files_path, img_name)
        label_path_root = os.path.join(label_files_path, img_name.split('.')[0])
        label_path = label_path_root + '.txt'
        img = Image.open(img_path)
        width = float(img.size[0])
        height = float(img.size[1])
        fo1 = open(label_path, 'r')
        lines = [l.split() for l in fo1.readlines() if l.strip()]
        S_uv = 0
        for line in lines:
            if line[0] == '2':
                # 像素数累加
                left = int(round(((float(line[1]) - 0.5*float(line[3]))*width), 0))
                top = int(round(((float(line[2]) - 0.5*float(line[4]))*height), 0))
                right = int(round(((float(line[1]) + 0.5*float(line[3]))*width), 0))
                low = int(round(((float(line[2]) + 0.5*float(line[4]))*height), 0))
                # print(left)
                # print(top)
                # print(right)
                # print(low)
                add_S_uv = num_uv(img_path, left, top, right, low)
                S_uv = S_uv + add_S_uv
            else: continue
        S_e = S_uv*f_dx*f_dx
        S_e = str(S_e) + ' mm^2'
        result.append([img_name, S_e])

    return result

def num_uv(img_path, left, top, right, low):
    mat_img = cv2.imread(img_path)
    cropImg = mat_img[top:low, left:right]
    # cv2.imshow("test", cropImg)
    # cv2.imwrite("./koutu/koutuqian/former_{}.jpg".format(num_jpg), cropImg)

    # 提取中心点的像素rgb值
    x_center = int(round(0.5*(right - left), 0))
    y_center = int(round(0.5*(low - top), 0))
    rgb = cropImg[y_center][x_center]
    # print(rgb)
    rgb_yuzhi = [x + 50 for x in rgb]
    # print(rgb_yuzhi)
    num_uv_ = 0
    for i in range(cropImg.shape[0]):
        for j in range(cropImg.shape[1]):
            # print(cropImg[i][j][0])
            # print(rgb_yuzhi[0])
            # print(cropImg[i][j])
            # print(rgb_yuzhi)
            if (cropImg[i][j][0] < rgb_yuzhi[0])&(cropImg[i][j][1] < rgb_yuzhi[1])&(cropImg[i][j][2] < rgb_yuzhi[2]):
                num_uv_ = num_uv_ + 1
            else: cropImg[i][j] = [1, 1, 1]
    # cv2.imwrite("./koutu/koutuhou/former_{}.jpg".format(num_jpg), cropImg)

    print("num_uv_")
    print(num_uv_)
    print("juxing")
    print((low-top)*(right-left))

    return num_uv_





img_files_path_5 = "./runs/detect/exp_0.5m_new/images"
label_files_path_5 = "./runs/detect/exp_0.5m_new/labels"
img_files_path_1 = "./runs/detect/exp_1m_new/images"
label_files_path_1 = "./runs/detect/exp_1m_new/labels"
f_dx_5 = 0.1283
f_dx_1 = 0.1867

result_5 = rust_estimation(img_files_path_5, label_files_path_5, f_dx_5)
# print(result_5)
# for line in result_5:
#     # print(line)
#     with open("result_0.5m", 'a') as f:
#         f.writelines(line[0] + ': ' + line[1] +'\n')
#     f.close()
#
#
# result_1 = rust_estimation(img_files_path_1, label_files_path_1, f_dx_1)
# print(result_1)
# for line in result_1:
#     # print(line)
#     with open("result_1m", 'a') as f:
#         f.writelines(line[0] + ': ' + line[1] +'\n')
#     f.close()